{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tce3stUlHN0L"
      },
      "source": [
        "##### Copyright 2019 The TensorFlow Authors.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "tuOe1ymfHZPu"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MfBg1C5NB3X0"
      },
      "source": [
        "# Keras による分散型トレーニング"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r6P32iYYV27b"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "<img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\"><a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs-l10n/blob/master/site/ja/tutorials/distribute/keras.ipynb\">TensorFlow.org で表示</a>\n",
        "</td>\n",
        "  <td>\n",
        "<img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\"><a target=\"_blank\" href=\"https://github.com/tensorflow/docs-l10n/blob/master/site/ja/tutorials/distribute/keras.ipynb\">Google Colab で実行</a>\n",
        "</td>\n",
        "  <td>\n",
        "<img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\"><a target=\"_blank\" href=\"https://storage.googleapis.com/tensorflow_docs/docs-l10n/site/ja/tutorials/distribute/keras.ipynb\">GitHub でソースを表示</a>\n",
        "</td>\n",
        "  <td>\n",
        "<img src=\"https://www.tensorflow.org/images/download_logo_32px.png\"><a href=\"https://storage.googleapis.com/tensorflow_docs/docs/site/en/tutorials/distribute/keras.ipynb\">ノートブックをダウンロード</a>\n",
        "</td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xHxb-dlhMIzW"
      },
      "source": [
        "## 概要\n",
        "\n",
        "`tf.distribute.Strategy` API は、複数の処理ユニットに渡ってトレーニングを分散するための抽象化を提供します。目標は、ユーザーが既存のモデルとトレーニングコードを使用して、最小限の変更で分散型トレーニングをできるようにすることです。\n",
        "\n",
        "このチュートリアルでは、`tf.distribute.MirroredStrategy`を使用します。これは、1 台のマシン上の多数の GPU で同期トレーニングを行うグラフ内レプリケーションを行います。基本的には、モデルの変数の全ての各プロセッサにコピーします。その後、[all-reduce](http://mpitutorial.com/tutorials/mpi-reduce-and-allreduce/)を使用して全プロセッサからの勾配を結合することで、結合された値をモデルの全コピーに適用します。\n",
        "\n",
        "`MirroredStrategy`は、TensorFlow コアで利用可能な数ある分散ストラテジーの中の 1 つです。その他のストラテジーについては、[分散ストラテジーガイド](../../guide/distributed_training.ipynb)をご覧ください。\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MUXex9ctTuDB"
      },
      "source": [
        "### Keras API\n",
        "\n",
        "この例では、`tf.keras` API を使用してモデルとトレーニングループを構築します。カスタムトレーニングループについては、[トレーニングループを使用した tf.distribute.Strategy](training_loops.ipynb) チュートリアルをご覧ください。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Dney9v7BsJij"
      },
      "source": [
        "## 依存関係をインポートする"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "r8S3ublR7Ay8"
      },
      "outputs": [],
      "source": [
        "# Import TensorFlow and TensorFlow Datasets\n",
        "\n",
        "import tensorflow_datasets as tfds\n",
        "import tensorflow as tf\n",
        "tfds.disable_progress_bar()\n",
        "\n",
        "import os"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SkocY8tgRd3H"
      },
      "outputs": [],
      "source": [
        "print(tf.__version__)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hXhefksNKk2I"
      },
      "source": [
        "## データセットをダウンロードする"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OtnnUwvmB3X5"
      },
      "source": [
        "MNIST データセットをダウンロードして [TensorFlow Datasets](https://www.tensorflow.org/datasets) から 読み込みます。これは`tf.data`形式のデータセットを返します。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lHAPqG8MtS8M"
      },
      "source": [
        "`with_info`を`True`に設定すると、データセット全体に対するメタデータを含み、ここでは`info`に保存されます。このメタデータオブジェクトは、トレーニングとテストの例の数を含みます。\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iXMJ3G9NB3X6"
      },
      "outputs": [],
      "source": [
        "datasets, info = tfds.load(name='mnist', with_info=True, as_supervised=True)\n",
        "\n",
        "mnist_train, mnist_test = datasets['train'], datasets['test']"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GrjVhv-eKuHD"
      },
      "source": [
        "## 分散ストラテジーを定義する"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TlH8vx6BB3X9"
      },
      "source": [
        "`MirroredStrategy`オブジェクトを作成します。これは分散を処理し、コンテキストマネージャ（`tf.distribute.MirroredStrategy.scope`）を提供して内側にモデルを構築します。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4j0tdf4YB3X9"
      },
      "outputs": [],
      "source": [
        "strategy = tf.distribute.MirroredStrategy()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cY3KA_h2iVfN"
      },
      "outputs": [],
      "source": [
        "print('Number of devices: {}'.format(strategy.num_replicas_in_sync))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lNbPv0yAleW8"
      },
      "source": [
        "## 入力パイプラインをセットアップする"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "psozqcuptXhK"
      },
      "source": [
        "マルチ GPU でモデルをトレーニングする場合、バッチサイズを増やして追加の計算能力を効果的に利用することができます。一般的には、GPU メモリに収まる最大のバッチサイズを使用し、それに応じて学習率を調整します。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "p1xWxKcnhar9"
      },
      "outputs": [],
      "source": [
        "# You can also do info.splits.total_num_examples to get the total\n",
        "# number of examples in the dataset.\n",
        "\n",
        "num_train_examples = info.splits['train'].num_examples\n",
        "num_test_examples = info.splits['test'].num_examples\n",
        "\n",
        "BUFFER_SIZE = 10000\n",
        "\n",
        "BATCH_SIZE_PER_REPLICA = 64\n",
        "BATCH_SIZE = BATCH_SIZE_PER_REPLICA * strategy.num_replicas_in_sync"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0Wm5rsL2KoDF"
      },
      "source": [
        "0-255 のピクセル値は[ 0-1 の範囲に正規化する必要があります](https://en.wikipedia.org/wiki/Feature_scaling)。このスケールを関数で定義します。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Eo9a46ZeJCkm"
      },
      "outputs": [],
      "source": [
        "def scale(image, label):\n",
        "  image = tf.cast(image, tf.float32)\n",
        "  image /= 255\n",
        "\n",
        "  return image, label"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WZCa5RLc5A91"
      },
      "source": [
        "この関数をトレーニングデータとテストデータに適用し、トレーニングデータをシャッフルし、[それトレーニング用にバッチ処理します](https://www.tensorflow.org/api_docs/python/tf/data/Dataset#batch)。パフォーマンスを向上のために、トレーニングデータのインメモリキャッシュも保持していることに注意してください。\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gRZu2maChwdT"
      },
      "outputs": [],
      "source": [
        "train_dataset = mnist_train.map(scale).cache().shuffle(BUFFER_SIZE).batch(BATCH_SIZE)\n",
        "eval_dataset = mnist_test.map(scale).batch(BATCH_SIZE)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4xsComp8Kz5H"
      },
      "source": [
        "## モデルを作成する"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1BnQYQTpB3YA"
      },
      "source": [
        "`strategy.scope`のコンテキストで Keras モデルを作成し、コンパイルします。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IexhL_vIB3YA"
      },
      "outputs": [],
      "source": [
        "with strategy.scope():\n",
        "  model = tf.keras.Sequential([\n",
        "      tf.keras.layers.Conv2D(32, 3, activation='relu', input_shape=(28, 28, 1)),\n",
        "      tf.keras.layers.MaxPooling2D(),\n",
        "      tf.keras.layers.Flatten(),\n",
        "      tf.keras.layers.Dense(64, activation='relu'),\n",
        "      tf.keras.layers.Dense(10)\n",
        "  ])\n",
        "\n",
        "  model.compile(loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
        "                optimizer=tf.keras.optimizers.Adam(),\n",
        "                metrics=['accuracy'])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8i6OU5W9Vy2u"
      },
      "source": [
        "## コールバックを定義する\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YOXO5nvvK3US"
      },
      "source": [
        "ここでは以下のコールバックを使用します。\n",
        "\n",
        "- *TensorBoard*: このコールバックはグラフの可視化を可能にする TensorBoard のためのログを書き込みます。\n",
        "- *モデルチェックポイント*<br>: このコールバックは各エポック後のモデルを保存します。\n",
        "- *学習率スケジューラ*: このコールバックを使用すると、各エポック/バッチ後に変化する学習率をスケジュールすることができます。\n",
        "\n",
        "例証目的として、このノートブックでは出力コールバックを追加して*学習率*を表示します。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "A9bwLCcXzSgy"
      },
      "outputs": [],
      "source": [
        "# Define the checkpoint directory to store the checkpoints\n",
        "\n",
        "checkpoint_dir = './training_checkpoints'\n",
        "# Name of the checkpoint files\n",
        "checkpoint_prefix = os.path.join(checkpoint_dir, \"ckpt_{epoch}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wpU-BEdzJDbK"
      },
      "outputs": [],
      "source": [
        "# Function for decaying the learning rate.\n",
        "# You can define any decay function you need.\n",
        "def decay(epoch):\n",
        "  if epoch &lt; 3:\n",
        "    return 1e-3\n",
        "  elif epoch &gt;= 3 and epoch &lt; 7:\n",
        "    return 1e-4\n",
        "  else:\n",
        "    return 1e-5"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jKhiMgXtKq2w"
      },
      "outputs": [],
      "source": [
        "# Callback for printing the LR at the end of each epoch.\n",
        "class PrintLR(tf.keras.callbacks.Callback):\n",
        "  def on_epoch_end(self, epoch, logs=None):\n",
        "    print('\\nLearning rate for epoch {} is {}'.format(epoch + 1,\n",
        "                                                      model.optimizer.lr.numpy()))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YVqAbR6YyNQh"
      },
      "outputs": [],
      "source": [
        "callbacks = [\n",
        "    tf.keras.callbacks.TensorBoard(log_dir='./logs'),\n",
        "    tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_prefix,\n",
        "                                       save_weights_only=True),\n",
        "    tf.keras.callbacks.LearningRateScheduler(decay),\n",
        "    PrintLR()\n",
        "]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "70HXgDQmK46q"
      },
      "source": [
        "## トレーニングして評価する"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6EophnOAB3YD"
      },
      "source": [
        "では、モデル上で`fit`を呼び出し、チュートリアルの最初に作成したデータセットを渡す、通常の方法でモデルをトレーニングします。トレーニングの分散の有無に関わらず、このステップは同じです。\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7MVw_6CqB3YD"
      },
      "outputs": [],
      "source": [
        "model.fit(train_dataset, epochs=12, callbacks=callbacks)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NUcWAUUupIvG"
      },
      "source": [
        "以下でご覧いただけるように、チェックポイントは保存されています。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JQ4zeSTxKEhB"
      },
      "outputs": [],
      "source": [
        "# check the checkpoint directory\n",
        "!ls {checkpoint_dir}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qor53h7FpMke"
      },
      "source": [
        "モデルのパフォーマンスを確認するには、最新のチェックポイントを読み込み、テストデータ上で`evaluate`を呼び出します。\n",
        "\n",
        "適切なデータセットを使用して、前と同様に`evaluate`を呼び出します。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JtEwxiTgpQoP"
      },
      "outputs": [],
      "source": [
        "model.load_weights(tf.train.latest_checkpoint(checkpoint_dir))\n",
        "\n",
        "eval_loss, eval_acc = model.evaluate(eval_dataset)\n",
        "\n",
        "print('Eval loss: {}, Eval Accuracy: {}'.format(eval_loss, eval_acc))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IIeF2RWfYu4N"
      },
      "source": [
        "出力の確認には、TensorBoard のログをダウンロードして端末に表示することができます。\n",
        "\n",
        "```\n",
        "$ tensorboard --logdir=path/to/log-directory\n",
        "```"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LnyscOkvKKBR"
      },
      "outputs": [],
      "source": [
        "!ls -sh ./logs"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kBLlogrDvMgg"
      },
      "source": [
        "## SavedModel にエクスポートする"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Xa87y_A0vRma"
      },
      "source": [
        "グラフと変数をプラットフォームに依存しない SavedModel 形式にエクスポートします。モデルが保存された後、スコープの有無に関わらずそれを読み込むことができます。\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "h8Q4MKOLwG7K"
      },
      "outputs": [],
      "source": [
        "path = 'saved_model/'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4HvcDmVsvQoa"
      },
      "outputs": [],
      "source": [
        "model.save(path, save_format='tf')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vKJT4w5JwVPI"
      },
      "source": [
        "`strategy.scope`なしでモデルを読み込みます。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "T_gT0RbRvQ3o"
      },
      "outputs": [],
      "source": [
        "unreplicated_model = tf.keras.models.load_model(path)\n",
        "\n",
        "unreplicated_model.compile(\n",
        "    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
        "    optimizer=tf.keras.optimizers.Adam(),\n",
        "    metrics=['accuracy'])\n",
        "\n",
        "eval_loss, eval_acc = unreplicated_model.evaluate(eval_dataset)\n",
        "\n",
        "print('Eval loss: {}, Eval Accuracy: {}'.format(eval_loss, eval_acc))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YBLzcRF0wbDe"
      },
      "source": [
        "`strategy.scope`と共にモデルを読み込みます。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BBVo3WGGwd9a"
      },
      "outputs": [],
      "source": [
        "with strategy.scope():\n",
        "  replicated_model = tf.keras.models.load_model(path)\n",
        "  replicated_model.compile(loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
        "                           optimizer=tf.keras.optimizers.Adam(),\n",
        "                           metrics=['accuracy'])\n",
        "\n",
        "  eval_loss, eval_acc = replicated_model.evaluate(eval_dataset)\n",
        "  print ('Eval loss: {}, Eval Accuracy: {}'.format(eval_loss, eval_acc))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MUZwaz4AKjtD"
      },
      "source": [
        "### 例とチュートリアル\n",
        "\n",
        "ここに keras fit/compile を用いた分散ストラテジーの使用例を紹介します。\n",
        "\n",
        "1. `tf.distribute.MirroredStrategy`を使用してトレーニングされる [Transformer](https://github.com/tensorflow/models/blob/master/official/nlp/transformer/transformer_main.py) の例。\n",
        "2. `tf.distribute.MirroredStrategy`を使用してトレーニングされる [NCF](https://github.com/tensorflow/models/blob/master/official/recommendation/ncf_keras_main.py) の例。\n",
        "\n",
        "[分散ストラテジーガイド](../../guide/distributed_training.ipynb#examples_and_tutorials)には他の例も多く記載されています。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8uNqWRdDMl5S"
      },
      "source": [
        "## 次のステップ\n",
        "\n",
        "- [分散ストラテジーガイド](../../guide/distributed_training.ipynb)を読みましょう。\n",
        "- [カスタムトレーニングループを使用した分散型トレーニング](training_loops.ipynb)チュートリアルを読みましょう。\n",
        "- 他のストラテジーや独自の TensorFlow モデルのパフォーマンス最適化に使用できる[ツール](../../guide/function.ipynb)についての詳細は、ガイドの[パフォーマンスのセクション](../../guide/profiler.md)をご覧ください。\n",
        "\n",
        "注意: `tf.distribute.Strategy`の開発は積極的に進められています。近日中にはより多くの例やチュートリアルを追加する予定なので、ぜひお試しください。フィードバックは [GitHub の Issue](https://github.com/tensorflow/tensorflow/issues/new) からお寄せください。"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "keras.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
